package icmp

import (
	"ConvertFlow/message"
	"fmt"
	"github.com/sirupsen/logrus"
	"log"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

const (
	ProtocolICMP = 1
)

// IcmpSendRaw Send ICMP echo request packet (code 8) to remote
func IcmpSendRaw(listeningReplyAddr string, remoteAddr string, data string) (*net.IPAddr, error) {
	// Listen用于获取PacketConn但我们不会等待回复
	c, err := icmp.ListenPacket("ip4:icmp", listeningReplyAddr)
	if err != nil {
		return nil, err
	}

	defer func(c *icmp.PacketConn) {
		err := c.Close()
		if err != nil {
			logrus.Println("icmp Close function failed")
		}
	}(c)
	// Resolve any DNS (if used) and get the real IP of the target
	dst, err := net.ResolveIPAddr("ip4", remoteAddr)
	if err != nil {
		return nil, err
	}

	// 构造ICMP 消息
	m := icmp.Message{
		Type: ipv4.ICMPTypeEcho,
		Code: 0,
		Body: &icmp.Echo{
			ID:   os.Getpid() & 0xffff,
			Seq:  1, //<< uint(seq), // TODO:正常情况下，序列号要进行每次递增才行
			Data: []byte(data),
		},
	}
	//从结构体转为字符数组
	b, err := m.Marshal(nil)
	if err != nil {
		return dst, err
	}

	// 将字符数据中内容发送出去
	n, err := c.WriteTo(b, dst)
	if err != nil {
		return dst, err
	} else if n != len(b) {
		return dst, fmt.Errorf("got %v; want %v", n, len(b))
	}

	return dst, nil
}

// Chunks Return a slice of a the string chunked with specific length
func Chunks(s string, chunkSize int) []string {
	if len(s) == 0 {
		return nil
	}
	if chunkSize >= len(s) {
		return []string{s}
	}
	var chunks = make([]string, 0, (len(s)-1)/chunkSize+1)
	currentLen := 0
	currentStart := 0
	for i := range s {
		if currentLen == chunkSize {
			chunks = append(chunks, s[currentStart:i])
			currentLen = 0
			currentStart = i
		}
		currentLen++
	}
	chunks = append(chunks, s[currentStart:])
	return chunks
}

// SendNoReply Send string to remote using ICMP
func SendNoReply(listeningReplyAddr string, remoteAddr string, chunkSize int, delay int, data string, savePcap bool) error {
	var w *pcapgo.Writer
	var f *os.File

	if savePcap {
		// 创建pcap文件名：源IP_目标IP_时间戳_数据长度.pcap
		timestamp := time.Now().Format("20060102150405")
		pcapFilename := fmt.Sprintf("%s_%s_%s_%d.pcap",
			listeningReplyAddr,
			remoteAddr,
			timestamp,
			len(data))

		// 创建pcap文件和写入器
		var err error
		f, err = os.Create(pcapFilename)
		if err != nil {
			return fmt.Errorf("create pcap file failed: %v", err)
		}
		defer f.Close()

		w = pcapgo.NewWriter(f)
		err = w.WriteFileHeader(65536, layers.LinkTypeIPv4)
		if err != nil {
			return fmt.Errorf("write pcap header failed: %v", err)
		}
	}

	// 解析IP地址
	dstIP := net.ParseIP(remoteAddr)
	if dstIP == nil {
		ips, err := net.LookupIP(remoteAddr)
		if err != nil || len(ips) == 0 {
			return fmt.Errorf("invalid IP address or hostname: %v", err)
		}
		dstIP = ips[0]
	}

	srcIP := net.ParseIP(listeningReplyAddr)
	if srcIP == nil || srcIP.IsUnspecified() {
		conn, err := net.Dial("udp", "8.8.8.8:80")
		if err != nil {
			return fmt.Errorf("get local IP failed: %v", err)
		}
		defer conn.Close()
		srcIP = conn.LocalAddr().(*net.UDPAddr).IP
	}

	dataSlice := Chunks(data, chunkSize)
	dataSlice = message.Marshall(dataSlice)
	time.Sleep(time.Duration(delay) * time.Second)

	// 发送数据包数量并记录到pcap
	countData := strconv.Itoa(len(dataSlice))
	if savePcap {
		if err := saveAndSendICMP(w, srcIP, dstIP, listeningReplyAddr, remoteAddr, 1, []byte(countData)); err != nil {
			return err
		}
	} else {
		if _, err := IcmpSendRaw(listeningReplyAddr, remoteAddr, countData); err != nil {
			return fmt.Errorf("send count packet failed: %v", err)
		}
	}

	// 发送实际数据包并记录到pcap
	for i := 0; i < len(dataSlice); i++ {
		time.Sleep(time.Duration(delay) * time.Second)
		if savePcap {
			if err := saveAndSendICMP(w, srcIP, dstIP, listeningReplyAddr, remoteAddr, i+2, []byte(dataSlice[i])); err != nil {
				return err
			}
		} else {
			if _, err := IcmpSendRaw(listeningReplyAddr, remoteAddr, dataSlice[i]); err != nil {
				return fmt.Errorf("send data packet %d failed: %v", i, err)
			}
		}
	}

	return nil
}

func SendWhileNoEchoReply(listeningReplyAddr string, remoteAddr string, data string) {
	for { //while we do not received echo reply ~ACK, resend it
		dst, dur, err := IcmpSendAndWaitForReply(listeningReplyAddr, remoteAddr, data)
		if err != nil {
			fmt.Println(err)
			fmt.Println("Retrying...")
		} else {
			fmt.Printf("PING %s (%s): %s\n", remoteAddr, dst, dur)
			break
		}
	}
}

func SendReply(listeningReplyAddr string, remoteAddr string, chunkSize int, delay int, data string) {
	dataSlice := Chunks(data, chunkSize) //1 character = 1byte , max size of icmp data 65507
	dataSlice = message.Marshall(dataSlice)
	time.Sleep(time.Duration(delay) * time.Second)

	// Announce the data size
	SendWhileNoEchoReply(listeningReplyAddr, remoteAddr, strconv.Itoa(len(dataSlice)))

	//Send the data
	for i := 0; i < len(dataSlice); i++ {
		time.Sleep(time.Duration(delay) * time.Second)
		SendWhileNoEchoReply(listeningReplyAddr, remoteAddr, dataSlice[i])
	}
}

func IcmpSendAndWaitForReply(listeningReplyAddr string, remoteAddr string, data string) (*net.IPAddr, time.Duration, error) {
	// Start listening for icmp replies
	c, err := icmp.ListenPacket("ip4:icmp", listeningReplyAddr)
	if err != nil {
		return nil, 0, err
	}
	defer c.Close()

	start := time.Now()
	dst, err := IcmpSendRaw(listeningReplyAddr, remoteAddr, data)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// Wait for a reply
	reply := make([]byte, 65535)
	err = c.SetReadDeadline(time.Now().Add(10 * time.Second))
	if err != nil {
		return dst, 0, err
	}
	n, peer, err := c.ReadFrom(reply) //n, peer, err := c.ReadFrom(reply)
	if err != nil {
		return dst, 0, err
	}
	duration := time.Since(start)

	// Pack it up boys, we're done here
	rm, err := icmp.ParseMessage(ProtocolICMP, reply[:n])
	if err != nil {
		return dst, 0, err
	}

	//We do not received echo reply
	if rm.Type != ipv4.ICMPTypeEchoReply {
		return dst, 0, fmt.Errorf("got %+v from %v; want echo reply", rm, peer)
	}
	return dst, duration, nil
}

// SendHashedmessage Send hashed message using ICMP
func SendHashedmessage(msg string, remoteAddr string, listenAddr string) {
	fmt.Println("addr", remoteAddr)
	// hash := utils.Sha1(msg)
	// fmt.Println("hash", hash)
	dst, dur, err := IcmpSendAndWaitForReply(listenAddr, msg, remoteAddr)
	if err != nil {
		panic(err)
	}
	log.Printf("PING %s (%s): %s\n", remoteAddr, dst, dur)
}

// saveAndSendICMP 保存ICMP数据包到pcap文件并实际发送
func saveAndSendICMP(w *pcapgo.Writer, srcIP, dstIP net.IP, listeningReplyAddr, remoteAddr string, seq int, data []byte) error {
	// 构造IPv4层
	ipLayer := &layers.IPv4{
		SrcIP:    srcIP,
		DstIP:    dstIP,
		Protocol: layers.IPProtocolICMPv4,
		Version:  4,
		TTL:      64,
		Id:       uint16(os.Getpid() & 0xffff),
	}

	// 构造ICMP层
	icmpLayer := &layers.ICMPv4{
		TypeCode: layers.CreateICMPv4TypeCode(layers.ICMPv4TypeEchoRequest, 0),
		Id:       uint16(os.Getpid() & 0xffff),
		Seq:      uint16(seq),
	}

	// 序列化数据包
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	payload := gopacket.Payload(data)

	err := gopacket.SerializeLayers(buf, opts,
		ipLayer,
		icmpLayer,
		payload,
	)
	if err != nil {
		return fmt.Errorf("serialize packet failed: %v", err)
	}

	// 写入pcap文件
	err = w.WritePacket(gopacket.CaptureInfo{
		Timestamp:     time.Now(),
		CaptureLength: len(buf.Bytes()),
		Length:        len(buf.Bytes()),
	}, buf.Bytes())
	if err != nil {
		return fmt.Errorf("write packet to pcap failed: %v", err)
	}

	// 实际发送ICMP包
	_, err = IcmpSendRaw(listeningReplyAddr, remoteAddr, string(data))
	if err != nil {
		return fmt.Errorf("send ICMP packet failed: %v", err)
	}

	return nil
}

// saveICMPPacket 构造并保存单个ICMP数据包
func saveICMPPacket(w *pcapgo.Writer, srcIP, dstIP net.IP, seq int, data []byte) error {
	// 构造以太网层
	eth := layers.IPv4{
		SrcIP:    srcIP,
		DstIP:    dstIP,
		Protocol: layers.IPProtocolICMPv4,
		Version:  4,
		TTL:      64,
		Id:       uint16(os.Getpid() & 0xffff),
	}

	// 构造ICMP层
	icmp := layers.ICMPv4{
		TypeCode: layers.CreateICMPv4TypeCode(layers.ICMPv4TypeEchoRequest, 0),
		Id:       uint16(os.Getpid() & 0xffff),
		Seq:      uint16(seq),
	}

	// 序列化数据包
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	// 设置ICMP负载数据
	payload := gopacket.Payload(data)

	err := gopacket.SerializeLayers(buf, opts,
		&eth,
		&icmp,
		payload,
	)
	if err != nil {
		return fmt.Errorf("serialize packet failed: %v", err)
	}

	// 写入pcap文件
	err = w.WritePacket(gopacket.CaptureInfo{
		Timestamp:     time.Now(),
		CaptureLength: len(buf.Bytes()),
		Length:        len(buf.Bytes()),
	}, buf.Bytes())
	if err != nil {
		return fmt.Errorf("write packet failed: %v", err)
	}

	return nil
}
